# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment3/question2/models.py
import torch.nn


# noinspection SpellCheckingInspection
class DebyeModel(torch.nn.Module):
    # 此默认值是由 main_guess.py 给出的
    def __init__(self,
                 epsilon_infinity: float = 1,
                 delta_epsilon: float = 5,
                 tau: float = 0):
        super().__init__()
        self.epsilon_infinity = torch.nn.Parameter(
            torch.tensor(epsilon_infinity, dtype=torch.double))
        self.delta_epsilon = torch.nn.Parameter(
            torch.tensor(delta_epsilon, dtype=torch.double))
        self.tau = torch.nn.Parameter(
            torch.tensor(tau, dtype=torch.double))

    def forward(self, omega):
        epsilon = self.epsilon_infinity + self.delta_epsilon / (1 + 1j * omega * self.tau)
        return torch.stack([torch.real(epsilon), torch.imag(epsilon)], dim=-1)
